package libtool65.link;

import java.io.IOException;
import java.io.OutputStream;
import java.util.HashMap;
import java.util.List;

import libtool65.o65.ExportDefinition;
import libtool65.o65.Module;
import libtool65.o65.Relocation;
import libtool65.o65.RelocationTable;
import libtool65.vin.RelocationEntry;
import libtool65.vin.RelocationList;

public class Linker {
	private HashMap<String, ModuleDefinition> modules = new HashMap<String, ModuleDefinition>();
	private SymbolTable symbols = new SymbolTable();
	private int code_start = 0;
	private byte[] code_block;

	public void add(String name, Module module) {
		ModuleDefinition module_definition = new ModuleDefinition(name, module);
		modules.put(name, module_definition);
		symbols.addModule(module_definition, module.getReferences(), module.getExports());
	}

	public boolean hasUnresolvedReferences() {
		return symbols.hasUnresolvedReferences();
	}

	public int link(int base_address, int zp_base) throws LinkerException {
		int total_size = calcTotalCodeSize();
		initCode(base_address, total_size, zp_base);
		for (ModuleDefinition module : modules.values())
		{
			relocate(module, module.getTextAddress() - code_start, module.getModule().getTextRelocations());
			relocate(module, module.getDataAddress() - code_start, module.getModule().getDataRelocations());
		}
		return total_size;
	}

	protected void relocate(ModuleDefinition module, int current_offset,
			RelocationTable relocations) throws LinkerException {
		for (Relocation relocation : relocations.toList())
		{
			int relocate_offset = 0;
			int abs_address = -1;
			switch (relocation.getSegment())
			{
			case Relocation.SEGMENT_TEXT:
				relocate_offset = getTextOffset(module);
				break;
			case Relocation.SEGMENT_DATA:
				relocate_offset = getDataOffset(module);
				break;
			case Relocation.SEGMENT_BSS:
				relocate_offset = getBssOffset(module);
				break;
			case Relocation.SEGMENT_ZERO:
				relocate_offset = getZpOffset(module);
				break;
			case Relocation.SEGMENT_UNDEFINED:	// external reference
				abs_address = resolveReference(module.getModule().getImportLabel(relocation.getSegmentIndex()));
				break;
			default: 
				System.err.println("error: relocation segment in module "+module.getName()+" not supported:");
				relocation.dumpInfo(System.err);
				throw new LinkerException();
			}
			switch (relocation.getType())
			{
			case Relocation.TYPE_WORD:
			{
				int new_address = (abs_address >= 0 ? abs_address 
						: (code_block[current_offset+relocation.getAbsOffset()+1] << 8) + code_block[current_offset+relocation.getAbsOffset()] + relocate_offset);
				code_block[current_offset+relocation.getAbsOffset()] = (byte) (new_address & 0xff);
				code_block[current_offset+relocation.getAbsOffset()+1] = (byte) ((new_address >> 8) & 0xff);
				break;
			}
			case Relocation.TYPE_HIGH:
			{
				int new_address = (abs_address >= 0 ? abs_address 
						: (code_block[current_offset+relocation.getAbsOffset()] << 8) + relocation.getLowByte() + relocate_offset);
				relocation.setCurrentLow(new_address & 0xff);
				code_block[current_offset+relocation.getAbsOffset()] = (byte) ((new_address >> 8) & 0xff);
				break;
			}
			case Relocation.TYPE_LOW:
			{
				int new_address = (abs_address >= 0 ? abs_address 
						: code_block[current_offset+relocation.getAbsOffset()] + relocate_offset);
				code_block[current_offset+relocation.getAbsOffset()] = (byte) (new_address & 0xff);
				break;
			}
			default:
				System.err.println("error: relocation type in module "+module.getName()+" not supported:");
				relocation.dumpInfo(System.err);
				throw new LinkerException();
			}
		}
	}

	protected int getZpOffset(ModuleDefinition module) {
		return module.getZpAddress() - module.getModule().getZpBase();
	}

	protected int getBssOffset(ModuleDefinition module) {
		return module.getBssAddress() - module.getModule().getBssBase();
	}

	protected int getDataOffset(ModuleDefinition module) {
		return module.getDataAddress() - module.getModule().getDataBase();
	}

	protected int getTextOffset(ModuleDefinition module) {
		return module.getTextAddress() - module.getModule().getTextBase();
	}

	public int resolveReference(String label) {
		Symbol symbol = symbols.resolve(label);
		ExportDefinition export = symbol.getModuleDefinition().getModule().getExportDefinition(label);
		switch (export.getSegmentId())
		{
		case Relocation.SEGMENT_ABSOLUTE:
			return export.getValue();
		case Relocation.SEGMENT_TEXT:
			return getTextOffset(symbol.getModuleDefinition()) + export.getValue();
		case Relocation.SEGMENT_DATA:
			return getDataOffset(symbol.getModuleDefinition()) + export.getValue();
		case Relocation.SEGMENT_BSS:
			return getBssOffset(symbol.getModuleDefinition()) + export.getValue();
		case Relocation.SEGMENT_ZERO:
			return getZpOffset(symbol.getModuleDefinition()) + export.getValue();
		}
		return -1;
	}
	
	public boolean isAbsoluteReference(String label) {
		Symbol symbol = symbols.resolve(label);
		ExportDefinition export = symbol.getModuleDefinition().getModule().getExportDefinition(label);
		return export.getSegmentId() == Relocation.SEGMENT_ABSOLUTE;
	}

	protected void initCode(int base_address, int total_size, int zp_base) throws LinkerException {
		code_start = base_address;
		code_block = new byte [total_size];
		int bss_offset = total_size;
		int zp_offset = (zp_base < 0 ? 0 : zp_base);
		int current_code_offset = 0;
		for (ModuleDefinition module : modules.values())
		{
			module.setBaseAddress(base_address);
			bss_offset = module.calcBssOffset(bss_offset);
			zp_offset = module.calcZpOffset(zp_offset);
			base_address += module.getModule().getCodeSize();
			current_code_offset = module.copyCode(code_block, current_code_offset);
		}
		if ((zp_offset & 0xff) != zp_offset)
		{
			System.err.println("error: zeropage overflow.");
			throw new LinkerException();
		}
	}

	protected int calcTotalCodeSize() {
		int total_size = 0;
		for (ModuleDefinition module : modules.values())
		{
			total_size += module.getModule().getCodeSize();
		}
		return total_size;
	}

	public List<String> getUnresolvedReferences() {
		return symbols.getUnresolvedReferences();
	}

	public void writeBinary(OutputStream file, boolean write_address_first) throws IOException
	{
		if (write_address_first)
		{
			file.write(code_start & 0xFF);
			file.write((code_start >> 8) & 0xFF);
		}
		file.write(code_block);
	}

	public RelocationList getRelocationList() throws LinkerException {
		RelocationList relocations = new RelocationList();
		for (ModuleDefinition module : modules.values())
		{
			getRelocationEntries(relocations, module, module.getTextAddress() - code_start, module.getModule().getTextRelocations());
			getRelocationEntries(relocations, module, module.getDataAddress() - code_start, module.getModule().getDataRelocations());
		}
		return relocations;
	}

	private void getRelocationEntries(RelocationList relocation_list,
			ModuleDefinition module, int current_offset, RelocationTable relocations) throws LinkerException {
		for (Relocation relocation : relocations.toList())
		{
			RelocationEntry entry = new RelocationEntry();
			entry.setAddress(current_offset+relocation.getAbsOffset());
			switch (relocation.getSegment())
			{
			case Relocation.SEGMENT_TEXT:
			case Relocation.SEGMENT_DATA:
			case Relocation.SEGMENT_BSS:
				entry.setCodeSegment();
				break;
			case Relocation.SEGMENT_ZERO:
				entry.setZeroSegment();
				break;
			case Relocation.SEGMENT_UNDEFINED:	// external reference
			{
				String label = module.getModule().getImportLabel(relocation.getSegmentIndex());
				if (isAbsoluteReference(label))
					continue;
				int abs_address = resolveReference(label);
				if (abs_address < 0)
					entry.setExtern(relocation.getSegmentIndex());
				else
					entry.setCodeSegment();
				break;
			}
			default: 
				System.err.println("error: relocation segment in module "+module.getName()+" not supported:");
				relocation.dumpInfo(System.err);
				throw new LinkerException();
			}
			switch (relocation.getType())
			{
			case Relocation.TYPE_WORD:
				entry.setTypeWord();
				break;
			case Relocation.TYPE_HIGH:
				entry.setTypeHigh(relocation.getCurrentLow());
				break;
			case Relocation.TYPE_LOW:
				entry.setTypeLow();
				break;
			default:
				System.err.println("error: relocation type in module "+module.getName()+" not supported:");
				relocation.dumpInfo(System.err);
				throw new LinkerException();
			}
			relocation_list.add(entry);
		}
	}

	public int getTotalSize() {
		int total_size = 0;
		for (ModuleDefinition module : modules.values())
		{
			total_size += module.getModule().getTotalCodeSize();
		}
		return total_size;
	}

	public int getZPSize() {
		int total_size = 0;
		for (ModuleDefinition module : modules.values())
		{
			total_size += module.getModule().getZpSize();
		}
		return total_size;
	}
}
